#!/usr/bin/env python3
"""
FLUX新分类参数节点 - 全新的参数分类系统
将参数按功能分为依赖配置、训练配置和性能配置三个独立节点
"""

import os
import sys
from typing import Dict, Any
import torch

# 导入ComfyUI模块
import folder_paths

# 导入基类
from .base_lora_trainer_node import BaseLoRATrainerNode

# 导入统一定义的参数类型
from .lora_trainer_utils.constants import (
    FLUX_TRAINING_PARAMS,
    FLUX_DEPS_CHUNK,
    FLUX_TRAIN_CHUNK,
    FLUX_PERF_CHUNK,
    FLUX_NEW_PARAMS_CATEGORY
)

# 使用独立的日志系统，不受其他脚本影响
from .lora_trainer_utils.custom_logger import get_logger

# 为这个模块创建独立的logger
logger = get_logger("flux_new_params")

# --- 内部连接类型定义 ---
DEPS_CONFIG_TYPE = FLUX_DEPS_CHUNK
TRAINING_CONFIG_TYPE = FLUX_TRAIN_CHUNK
PERFORMANCE_CONFIG_TYPE = FLUX_PERF_CHUNK
# --- 最终输出类型定义 ---
FINAL_CONFIG_TYPE = FLUX_TRAINING_PARAMS


class FluxDependencyConfig:
    """FLUX依赖配置节点 - 管理训练所需的外部模型依赖路径"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # 模型路径配置
                "pretrained_model_name_or_path": (folder_paths.get_filename_list("checkpoints"), {
                    "default": "", 
                    "multiline": False, 
                    "tooltip": "预训练模型路径 - 选择要训练的基础模型。具体影响：SD1.5(通用模型，兼容性好，训练稳定) → SDXL(高质量模型，细节丰富，需要更多显存) → SDXL Turbo(快速模型，训练快但效果可能不如SDXL)。选择正确的模型对训练效果至关重要。"
                }),
                "clip_l": (["None"] + folder_paths.get_filename_list("clip"), {
                    "default": "None",
                    "tooltip": "CLIP-L模型文件 - FLUX.1训练必需的文本编码器，负责将文本转换为向量表示。CLIP-L是CLIP的改进版本，提供更好的文本理解能力。选择None时使用基础模型内置CLIP-L"
                }),
                "t5xxl": (["None"] + folder_paths.get_filename_list("text_encoders"), {
                    "default": "None",
                    "tooltip": "T5XXL模型文件 - FLUX.1训练必需的文本编码器，负责处理长文本序列。T5XXL基于T5架构，专门为FLUX模型优化，支持更长的文本输入和更丰富的语义理解。注意：此列表可能包含多个文件夹的模型，请选择包含't5'、't5xxl'、'flan-t5'等关键词的模型文件。选择None时使用基础模型内置T5XXL"
                }),
                "ae": (["None"] + folder_paths.get_filename_list("vae"), {
                    "default": "None",
                    "tooltip": "AutoEncoder模型文件 - FLUX.1训练必需的图像编码器/解码器，负责将图像转换为潜在空间表示和反向转换。AE模型决定了图像的质量和细节保留程度。选择None时使用基础模型内置AE"
                }),
                "pretrained_model_name_or_path": (folder_paths.get_filename_list("checkpoints"), {
                    "default": "", 
                    "multiline": False, 
                    "tooltip": "预训练模型路径 - 选择要训练的基础模型。具体影响：SD1.5(通用模型，兼容性好，训练稳定) → SDXL(高质量模型，细节丰富，需要更多显存) → SDXL Turbo(快速模型，训练快但效果可能不如SDXL)。选择正确的模型对训练效果至关重要。"
                }),
                # 基础训练参数
                "train_data_dir": ("STRING", {
                    "default": "", 
                    "multiline": False, 
                    "tooltip": "训练数据目录路径 - 包含训练图像的文件夹路径。具体要求：文件夹内应包含子文件夹，每个子文件夹包含图像文件和对应的文本描述文件(.txt或.caption)。例如：/path/to/train_data/character1/image1.jpg + image1.txt。注意：train_data_dir应该是包含子文件夹的父目录，而不是直接包含图像的目录。"
                }),
                "output_dir": ("STRING", {
                    "default": "./output", 
                    "multiline": False, 
                    "tooltip": "模型输出目录 - 训练完成的LoRA模型保存路径。具体影响：相对路径(相对于ComfyUI根目录) → 绝对路径(完整路径，推荐) → 默认路径(自动生成)。建议使用绝对路径，避免路径问题。"
                }),
                "output_name": ("STRING", {
                    "default": "trained_flux_lora", 
                    "multiline": False, 
                    "tooltip": "输出模型名称 - 训练完成的LoRA模型文件名(不含扩展名)。具体影响：简单名称(如'my_lora') → 描述性名称(如'character_style_v1') → 版本化名称(如'portrait_style_v2.1')。建议使用描述性名称，便于管理多个模型。"
                }),
                # 模型类型配置
                "config_file": ("STRING", {
                    "default": "", 
                    "multiline": False, 
                    "tooltip": "TOML配置文件路径 - 可选的配置文件路径，用于加载预定义的训练配置。具体用途：空值(使用节点参数) → TOML文件路径(加载配置文件，适合批量训练或复杂配置)。配置文件可以包含完整的训练参数设置，便于重复使用和版本管理。"
                }),
                # 内存优化配置
                "fp8_base": ("BOOLEAN", {
                    "default": True, 
                    "tooltip": "FP8基础模型 - 使用FP8精度加载基础模型，节省显存。具体效果：启用后显存占用减少30-50%，训练速度提升10-20%，是现代GPU的推荐设置。强烈建议启用，能显著提升训练效率。"
                }),
            }
        }

    RETURN_TYPES = (FLUX_TRAINING_PARAMS,)
    FUNCTION = "create_config"
    CATEGORY = FLUX_NEW_PARAMS_CATEGORY

    def validate_inputs(kwargs):
        """验证输入参数 - 返回True或错误字符串"""
 

        # 验证训练数据目录 - 只有当用户输入了值时才验证
        train_data_dir = kwargs.get("train_data_dir", "")
        if train_data_dir:  # 只有当用户输入了值时才验证
            if not os.path.exists(train_data_dir):
                return f"错误：训练数据目录不存在: {train_data_dir}"
            
            if not os.path.isdir(train_data_dir):
                return f"错误：训练数据路径不是目录: {train_data_dir}"
            
            # 检查是否有图片文件
            image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.webp']
            has_images = False
            try:
                for root, dirs, files in os.walk(train_data_dir):
                    for file in files:
                        if any(file.lower().endswith(ext) for ext in image_extensions):
                            has_images = True
                            break
                    if has_images:
                        break
                
                if not has_images:
                    return f"错误：训练数据目录中未找到图片文件: {train_data_dir}"
            except Exception as e:
                return f"错误：无法访问训练数据目录: {e}"
        
        # 验证输出目录 - 只有当用户输入了值时才验证
        output_dir = kwargs.get("output_dir", "")
        if output_dir:  # 只有当用户输入了值时才验证
            # 尝试创建输出目录
            try:
                os.makedirs(output_dir, exist_ok=True)
            except Exception as e:
                return f"错误：无法创建输出目录 {output_dir}: {e}"
        
        # 验证输出名称 - 只有当用户输入了值时才验证
        output_name = kwargs.get("output_name", "")
        if output_name:  # 只有当用户输入了值时才验证
            # 检查输出名称是否包含非法字符
            import re
            if re.search(r'[<>:"/\\|?*]', output_name):
                return f"错误：输出名称包含非法字符: {output_name}"
        
        # 验证配置文件路径（如果提供）
        config_file = kwargs.get("config_file", "")
        if config_file:
            if not os.path.exists(config_file):
                return f"错误：配置文件不存在: {config_file}"
            
            if not config_file.lower().endswith('.toml'):
                return f"错误：配置文件必须是TOML格式: {config_file}"
        
        # 验证必需模型文件 - 只有当用户选择了值时才验证
        clip_l = kwargs.get("clip_l", "")
        t5xxl = kwargs.get("t5xxl", "")
        ae = kwargs.get("ae", "")
        
        if clip_l and not os.path.exists(folder_paths.get_full_path("clip", clip_l)):
            return f"错误：CLIP-L模型文件不存在: {clip_l}"
        
        if t5xxl and not os.path.exists(folder_paths.get_full_path("text_encoders", t5xxl)):
            return f"错误：T5XXL模型文件不存在: {t5xxl}"
        
        if ae and not os.path.exists(folder_paths.get_full_path("vae", ae)):
            return f"错误：AutoEncoder模型文件不存在: {ae}"
        
        return True  # 验证通过

    def create_config(self, **kwargs):
        # 将相对路径转换为绝对路径
        for key in ["clip_l", "t5xxl", "ae","pretrained_model_name_or_path"]:
            if kwargs[key]:
                try:
                    if key == "clip_l":
                        kwargs[key] = folder_paths.get_full_path("clip", kwargs[key])
                    elif key == "t5xxl":
                        kwargs[key] = folder_paths.get_full_path("text_encoders", kwargs[key])
                    elif key == "ae":
                        kwargs[key] = folder_paths.get_full_path("vae", kwargs[key])
                    elif key == "pretrained_model_name_or_path":
                        kwargs[key] = folder_paths.get_full_path("checkpoints", kwargs[key])
                except Exception as e:
                    logger.warning(f"无法获取 {key} 的完整路径: {e}")
        
        return (kwargs,)


class FluxTrainingConfig:
    """FLUX训练配置节点 - 影响最终模型效果的参数"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # 网络结构参数 - 影响模型学习能力
                "network_dim": ("INT", {
                    "default": 32,  # 更新为更好的默认值
                    "min": 1, 
                    "max": 512, 
                    "tooltip": "LoRA网络维度(RANK) - 控制LoRA网络的表达能力。数值越大，模型学习能力越强，出图效果越好，但显存占用也越多。具体效果：2-8(基础效果，面部模糊，纹理粗糙) → 16-32(中等效果，面部清晰，有基本纹理) → 64-128(优秀效果，面部精致，纹理细腻，背景丰富)。建议值：16GB显存用32-64，24GB+显存用64-128。"
                }),
                "network_alpha": ("FLOAT", {
                    "default": 16.0,  # 更新为更好的默认值
                    "min": 0.1, 
                    "max": 128.0,  # 修复：从10.0改为128.0，与其他配置保持一致
                    "tooltip": "LoRA网络alpha值 - 控制LoRA权重更新的缩放因子。通常设置为network_dim的1/4到1/2，影响训练的稳定性和收敛速度。具体影响：较小值(0.1-0.5)训练稳定但学习慢，较大值(1.0-2.0)学习快但可能不稳定。推荐：network_dim的1/4到1/2，如network_dim=64时用16-32。"
                }),
                "network_weights": ("STRING", {
                    "default": "", 
                    "multiline": False, 
                    "tooltip": "网络权重路径 - 加载预训练的LoRA权重继续训练。具体用途：空值(从头开始训练) → 现有LoRA文件(继续训练，适合微调或增量训练)。加载现有权重可以加快收敛速度，但需要确保兼容性。"
                }),
                "network_train_unet_only": ("BOOLEAN", {
                    "default": True, 
                    "tooltip": "仅训练UNet - 只训练UNet部分，不训练文本编码器。具体影响：启用后训练速度提升20-30%，显存占用减少，但可能影响文本理解能力。建议：如果只关注图像生成质量，可以启用；如果需要更好的文本控制，保持关闭。"
                }),
                "network_train_text_encoder_only": ("BOOLEAN", {
                    "default": False, 
                    "tooltip": "仅训练文本编码器 - 只训练文本编码器部分。具体影响：启用后训练速度提升，显存占用减少，但可能影响图像生成质量。建议：通常保持关闭，除非有特殊的文本理解需求。"
                }),
                
                # 学习控制参数 - 影响训练效果
                "learning_rate": ("FLOAT", {
                    "default": 1e-4, 
                    "min": 1e-6, 
                    "max": 1e-2, 
                    "step": 1e-6, 
                    "tooltip": "学习率 - 控制模型学习速度的核心参数。具体出图效果：1e-5(学习很慢，细节学习充分但训练时间长) → 1e-4(标准学习率，平衡效果和速度，推荐) → 1e-3(学习很快，可能错过细节) → 1e-2(学习过快，容易过拟合)。1e-4通常能产生最好的训练效果，让模型既学到细节又不会过拟合。"
                }),
                "max_train_steps": ("INT", {
                    "default": 1200,  # 更新为更好的默认值
                    "min": 1, 
                    "max": 10000, 
                    "tooltip": "最大训练步数 - 控制训练的总时长。步数越多，模型学习越充分，但也更容易过拟合。具体效果：200步(基础学习，能学基本概念) → 500步(充分学习，效果较好) → 1000步(深入学习，效果优秀) → 2000步(过度学习，可能过拟合)。建议：小数据集(10-50张)200-500步，中等数据集(50-200张)500-1000步，大数据集(200+张)1000-2000步。"
                }),
                 # 保存与中断
                "save_every_n_steps": ("INT", {
                    "default": 500,
                    "min": 100,
                    "max": 5000,
                    "tooltip": "每多少步保存一次模型"
                }),
                "lr_scheduler": (["constant", "constant_with_warmup", "cosine", "cosine_with_restarts", "linear", "polynomial", "polynomial_with_warmup"], {
                    "default": "constant_with_warmup", 
                    "tooltip": "学习率调度器 - 控制学习率如何随时间变化。具体出图效果：constant(固定学习率，简单稳定) → constant_with_warmup(预热后固定，推荐，避免初期不稳定) → cosine(余弦衰减，学习效果最好) → linear(线性衰减，平衡效果) → polynomial(多项式衰减，适合长期训练)。constant_with_warmup通常能产生最稳定的训练效果。"
                }),
                "lr_warmup_steps": ("INT", {
                    "default": 100, 
                    "min": 0, 
                    "max": 1000, 
                    "tooltip": "学习率预热步数 - 让模型逐渐适应训练。具体影响：0(无预热，可能初期不稳定) → 100(标准预热，推荐) → 500(长预热，适合复杂模型) → 1000(超长预热，适合大型数据集)。预热能让模型在训练初期更稳定，避免学习率过高导致的震荡。"
                }),
                "lr_scheduler_num_cycles": ("INT", {
                    "default": 1, 
                    "min": 1, 
                    "max": 10, 
                    "tooltip": "学习率调度器周期数 - 控制学习率变化周期。具体影响：1(单周期，标准训练) → 2(双周期，适合长期训练) → 5(多周期，适合复杂任务)。周期数越多，学习率变化越频繁，可能提高训练效果但增加训练时间。"
                }),
                "max_grad_norm": ("FLOAT", {
                    "default": 0.0, 
                    "min": 0.0, 
                    "max": 10.0, 
                    "step": 0.1, 
                    "tooltip": "梯度裁剪阈值 - 防止梯度爆炸，稳定训练。具体影响：0.0(无裁剪，可能梯度爆炸) → 0.5(严格裁剪，训练稳定但可能学习不足) → 1.0(标准裁剪，平衡稳定性和学习效果，推荐) → 2.0(宽松裁剪，学习充分但可能不稳定) → 5.0(很少裁剪，可能梯度爆炸)。1.0通常能提供最佳的训练稳定性。"
                }),
                "weight_decay": ("FLOAT", {
                    "default": 0.01, 
                    "min": 0.0, 
                    "max": 0.1, 
                    "step": 0.001, 
                    "tooltip": "权重衰减 - 防止过拟合的正则化参数。具体出图效果：0.0(无衰减，可能过拟合) → 0.01(标准衰减，平衡效果和泛化，推荐) → 0.05(强衰减，泛化好但可能欠拟合) → 0.1(极强衰减，泛化最好但效果可能下降)。0.01通常能提供最佳的训练效果和泛化能力。"
                }),
                
                # 训练质量参数 - 影响最终效果
                "resolution": ("INT", {
                    "default": 512, 
                    "min": 256, 
                    "max": 2048, 
                    "step": 64, 
                    "tooltip": "训练分辨率 - 训练图像的处理尺寸。分辨率越高，细节保留越好，但显存占用呈平方增长。具体效果：256(低分辨率，细节模糊) → 512(标准分辨率，细节清晰) → 768(高分辨率，细节丰富) → 1024(超高分辨率，细节极致)。建议：16GB显存用512，24GB+显存可用768-1024。"
                }),
                "guidance_scale": ("FLOAT", {
                    "default": 1.0, 
                    "min": 0.1, 
                    "max": 10.0, 
                    "tooltip": "引导尺度 - FLUX.1训练中的引导强度参数。控制模型对条件信息的依赖程度，影响生成图像的质量和一致性。具体效果：0.1-0.5(弱引导，创意性强但可能偏离训练) → 1.0(标准引导，平衡效果) → 2.0-5.0(强引导，更接近训练数据) → 5.0+(极强引导，严格遵循训练)。建议：1.0-2.0，追求创意用0.5-1.0，追求准确性用2.0-5.0。"
                }),
                "timestep_sampling": (["shift", "uniform"], {
                    "default": "shift", 
                    "tooltip": "时间步采样方法 - 控制训练过程中时间步的选择策略。具体效果：shift(偏移采样，FLUX.1推荐，训练更稳定，效果更好) → uniform(均匀采样，更简单但可能效果较差，训练不稳定)。建议：使用shift，这是FLUX.1的官方推荐方法。"
                }),
                
                # 保存配置
                "save_every_n_epochs": ("INT", {
                    "default": 1, 
                    "min": 1, 
                    "max": 10, 
                    "tooltip": "每N轮保存一次模型 - 控制模型保存频率。具体影响：1(每轮保存，占用空间大但安全) → 2(每2轮保存，平衡空间和安全) → 5(每5轮保存，节省空间)。建议设置为1，可以随时恢复训练或选择最佳模型。"
                }),
                "save_model_as": (["safetensors", "ckpt", "diffusers"], {
                    "default": "safetensors", 
                    "tooltip": "模型保存格式 - 选择模型文件格式。具体特点：safetensors(安全格式，推荐，加载快，兼容性好) → ckpt(传统格式，兼容性最好) → diffusers(扩散器格式，适合特定用途)。safetensors是现代推荐格式，安全且高效。"
                }),
            }
        }
    
    RETURN_TYPES = (TRAINING_CONFIG_TYPE,)
    FUNCTION = "create_config"
    CATEGORY = FLUX_NEW_PARAMS_CATEGORY

    def create_config(self, **kwargs):
        # 处理"None"值，删除对应参数
        none_params = ["clip_l", "t5xxl", "ae"]
        for param in none_params:
            if kwargs.get(param) == "None":
                kwargs.pop(param, None)
                print(f"[INFO] 参数 {param} 设置为None，已从训练参数中移除")
        
        # 过滤空字符串参数 - 特别是network_weights
        filtered_kwargs = {}
        for key, value in kwargs.items():
            # 如果是字符串类型且为空，则跳过
            if isinstance(value, str) and not value.strip():
                continue
            filtered_kwargs[key] = value
        
        # 处理分辨率格式
        filtered_kwargs['resolution'] = f"{filtered_kwargs['resolution']},{filtered_kwargs['resolution']}"
        
        # 添加一些与训练效果相关的硬编码参数
        filtered_kwargs.update({
            "model_prediction_type": "raw",
            "discrete_flow_shift": 3.1582,
        })
        return (filtered_kwargs,)


class FluxPerformanceConfig:
    """FLUX性能配置节点 - 影响训练速度和显存占用的参数"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # 批次和梯度参数 - 影响显存占用
                "train_batch_size": ("INT", {
                    "default": 1, 
                    "min": 1, 
                    "max": 16, 
                    "tooltip": "训练批次大小 - 控制每次训练的样本数量。具体影响：1(标准训练，显存占用少) → 2(高效训练，学习效果更好) → 4(超高效训练，需要大量显存) → 8(极限训练，显存占用很大) → 16(超极限训练，需要顶级GPU)。批次越大训练效果越好，但显存占用呈线性增长。"
                }),
                "gradient_accumulation_steps": ("INT", {
                    "default": 1, 
                    "min": 1, 
                    "max": 16, 
                    "tooltip": "梯度累积步数 - 模拟更大批次训练的技术。具体影响：1(无累积，标准训练) → 2(等效批次×2，效果更好) → 4(等效批次×4，效果显著提升) → 8(等效批次×8，效果极佳) → 16(等效批次×16，效果最好但训练慢)。数值越大=等效批次越大=训练效果越好，但训练时间越长。"
                }),
                
                # 内存优化参数
                "blocks_to_swap": ("INT", {
                    "default": 8,  # 从12减少到8，提高训练速度
                    "min": 0, 
                    "max": 20, 
                    "tooltip": "交换到CPU的模型块数 - 精确控制显存和内存的平衡。具体影响：0(不交换，训练最快但显存占用最大) → 4(少量交换，平衡速度和显存) → 8(标准交换，推荐设置) → 12(大量交换，显存占用少但训练慢) → 16(超量交换，显存占用最少但训练很慢) → 20(极限交换，显存占用最少但训练极慢)。数值越大显存占用越少，但训练速度越慢。"
                }),
                "max_data_loader_n_workers": ("INT", {
                    "default": 8,  # 从4增加到8，进一步提高数据加载速度
                    "min": 0, 
                    "max": 8,  # 增加最大值限制
                    "tooltip": "数据加载器工作进程数，推荐值8(高速数据加载)"
                }),
                
                # 优化器配置
                "optimizer_type": (["AdamW", "AdamW8bit", "Lion", "SGDNesterov", "DAdaptation", "Prodigy", "AdaFactor"], {
                    "default": "Lion",  # 从AdamW8bit改为Lion，最快的优化器
                    "tooltip": "优化器类型 - 选择训练算法。具体效果：AdamW(标准优化器，效果最好，显存占用大) → AdamW8bit(8位AdamW，平衡效果和显存) → Lion(新优化器，速度快，显存占用中等) → Lion8bit(8位Lion，速度最快，显存占用最少) → SGDNesterov(传统优化器，显存占用少但效果可能不如AdamW) → SGDNesterov8bit(8位SGD，显存占用最少但效果最差) → AdaFactor(内存友好优化器，显存占用最少)。Lion8bit是速度最快的选择。"
                }),
                "mixed_precision": (["fp16", "bf16", "no"], {
                    "default": "bf16",  # 调整为与Flux16GBParams一致
                    "tooltip": "混合精度训练 - 控制计算精度。具体效果：fp16(16位浮点，兼容性好，显存节省30%，推荐) → bf16(脑浮点16位，现代GPU推荐，精度好，显存节省30%) → no(全精度32位，最稳定但显存占用大)。fp16是最稳定的选择，与gradient_checkpointing兼容性最好。"
                }),
                
                # 高级优化参数
                "highvram": ("BOOLEAN", {
                    "default": False, 
                    "tooltip": "高显存模式 - 启用更多显存优化选项。具体效果：启用后训练速度提升10-20%，能更好地利用显存资源，但需要足够的显存支持。建议：24GB+显存可以启用，能显著提升训练效率。"
                }),
                "full_bf16": ("BOOLEAN", {
                    "default": False, 
                    "tooltip": "全BF16模式 - 强制所有计算使用BF16精度。具体效果：启用后训练速度提升10-15%，但需要RTX 30/40系列或更新的GPU。现代GPU强烈推荐启用，能显著提升训练效率。"
                }),
                "sdpa": ("BOOLEAN", {
                    "default": True, 
                    "tooltip": "使用SDPA - 启用Scaled Dot-Product Attention优化。具体效果：启用后训练速度提升5-10%，能更好地处理注意力机制，提高训练效果。现代GPU强烈推荐启用。"
                }),
                
                # CPU卸载参数
                "enable_sequential_cpu_offload": ("BOOLEAN", {
                    "default": False, 
                    "tooltip": "启用顺序CPU卸载 - 将模型部分卸载到CPU以节省显存。具体影响：启用后显存占用减少，但训练速度显著下降(50-80%)。建议：只有显存严重不足时才启用，否则会大幅降低训练效率。"
                }),
                "enable_cpu_offload": ("BOOLEAN", {
                    "default": False, 
                    "tooltip": "启用CPU卸载 - 将部分模型卸载到CPU。具体影响：启用后显存占用减少，但训练速度下降(30-50%)。建议：显存不足时可以考虑，但会降低训练效率。"
                }),
                
                # 缓存策略
                "cache_latents": ("BOOLEAN", {
                    "default": True, 
                    "tooltip": "缓存Latents - 缓存图像编码结果。具体效果：启用后训练速度提升50-100%，能避免重复编码图像，大幅提升训练效率。强烈建议启用，是提升训练速度最有效的方法之一。"
                }),
                "cache_text_encoder_outputs": ("BOOLEAN", {
                    "default": True, 
                    "tooltip": "缓存文本编码器输出 - 缓存文本编码结果。具体效果：启用后训练速度提升30-50%，特别是处理长文本时效果明显，能避免重复计算文本编码。强烈建议启用。"
                }),
                "cache_text_encoder_outputs_to_disk": ("BOOLEAN", {
                    "default": True, 
                    "tooltip": "将文本编码器输出缓存到磁盘 - 节省内存。具体效果：启用后内存占用减少，训练速度提升，因为文本编码计算昂贵且结果可以重复使用。强烈建议启用。"
                }),
                
                # 分桶设置
                "enable_bucket": ("BOOLEAN", {
                    "default": True, 
                    "tooltip": "启用图像分桶 - 将相似尺寸图像分组训练，提高效率。具体效果：启用后训练速度提升20-30%，能更好地处理不同尺寸的图像，提高训练效果。强烈建议启用。"
                }),
                "min_bucket_reso": ("INT", {
                    "default": 256, 
                    "min": 128, 
                    "max": 512, 
                    "step": 64, 
                    "tooltip": "最小分桶分辨率 - 设置图像分桶的最小尺寸。具体影响：128(小图像分桶，适合细节训练) → 256(标准分桶，推荐) → 512(大图像分桶，适合整体训练)。256通常能提供最佳的训练效果。"
                }),
                "max_bucket_reso": ("INT", {
                    "default": 512, 
                    "min": 512, 
                    "max": 2048, 
                    "step": 64, 
                    "tooltip": "最大分桶分辨率 - 设置图像分桶的最大尺寸。具体影响：512(小图像分桶，训练快) → 1024(标准分桶，推荐) → 2048(大图像分桶，细节丰富)。1024通常能提供最佳的训练效果和速度平衡。"
                }),
                "bucket_reso_steps": ("INT", {
                    "default": 64, 
                    "min": 32, 
                    "max": 128, 
                    "step": 32, 
                    "tooltip": "分桶分辨率步长 - 控制分桶的精细程度。具体影响：32(精细分桶，训练效果最好但速度慢) → 64(标准分桶，推荐) → 128(粗糙分桶，训练快但效果可能下降)。64通常能提供最佳的效果和速度平衡。"
                }),
                
                # 训练控制参数
                "max_train_epochs": ("INT", {
                    "default": None, 
                    "min": 0, 
                    "max": 1000, 
                    "step": 1, 
                    "tooltip": "最大训练轮数 - None表示使用max_train_steps控制训练时长，大于0时优先使用轮数控制。具体影响：None(使用步数控制，推荐) → 1-10(使用轮数控制，适合固定轮数训练) → 100+(长期训练)。建议设置为None，让训练器使用max_train_steps来控制训练时长。"
                }),
            }
        }

    RETURN_TYPES = (PERFORMANCE_CONFIG_TYPE,)
    FUNCTION = "create_config"
    CATEGORY = FLUX_NEW_PARAMS_CATEGORY

    def create_config(self, **kwargs):
        # 将max_train_epochs为0时转换为None，这样训练器会使用max_train_steps
        if kwargs.get("max_train_epochs") == 0:
            kwargs["max_train_epochs"] = None
        
        return (kwargs,)


class FluxConfigCombiner:
    """FLUX配置合并节点 - 合并依赖、训练和性能配置"""

    aux_id = "comfyui_lora_train"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "deps_config": (FLUX_TRAINING_PARAMS, {
                    "tooltip": "依赖配置 - 包含CLIP-L、T5XXL、AE模型路径的配置块"
                }),
                "training_config": (TRAINING_CONFIG_TYPE, {
                    "tooltip": "训练配置 - 包含网络维度、学习率、训练步数等影响模型效果的参数"
                }),
                "performance_config": (PERFORMANCE_CONFIG_TYPE, {
                    "tooltip": "性能配置 - 包含批次大小、优化器、缓存设置等影响训练速度和显存占用的参数"
                }),
            }
        }

    RETURN_TYPES = (FINAL_CONFIG_TYPE,)
    FUNCTION = "combine_configs"
    CATEGORY = FLUX_NEW_PARAMS_CATEGORY

    def combine_configs(self, deps_config, training_config, performance_config):
        # 合并三个字典
        combined_config = {}
        combined_config.update(deps_config)
        combined_config.update(training_config)
        combined_config.update(performance_config)
        
        # 补全与Flux16GBParams一致的关键参数
        # 分辨率处理
        if 'resolution' in combined_config and isinstance(combined_config['resolution'], int):
            combined_config['resolution'] = f"{combined_config['resolution']},{combined_config['resolution']}"
        
        # 根据优化器类型设置相应的参数
        optimizer_type = combined_config.get("optimizer_type", "AdamW")
        if optimizer_type == "AdaFactor":
            combined_config["optimizer_args"] = ["relative_step=False", "scale_parameter=False", "warmup_init=False"]
        elif optimizer_type == "AdamW":
            combined_config["optimizer_args"] = ["weight_decay=0.01", "betas=0.9,0.999"]
        elif optimizer_type == "AdamW8bit":
            combined_config["optimizer_args"] = ["weight_decay=0.01", "betas=0.9,0.999"]
        elif optimizer_type == "Lion":
            combined_config["optimizer_args"] = ["weight_decay=0.01"]
        elif optimizer_type == "Lion8bit":
            combined_config["optimizer_args"] = ["weight_decay=0.01"]
        elif optimizer_type == "SGDNesterov":
            combined_config["optimizer_args"] = ["momentum=0.9", "weight_decay=0.01"]
        elif optimizer_type == "SGDNesterov8bit":
            combined_config["optimizer_args"] = ["momentum=0.9", "weight_decay=0.01"]
        else:
            # 默认AdamW参数
            combined_config["optimizer_args"] = ["weight_decay=0.01", "betas=0.9,0.999"]
        
        # 添加16GB节点的硬编码参数，但保持与Flux16GBParams一致
        fixed_params = {
            "network_module": "networks.lora_flux",
            "flux": True,
            "v2": False,
            "v_parameterization": False,
            "gradient_checkpointing": True,
            "xformers": False,
            "save_model_as": "safetensors",
            "persistent_data_loader_workers": True,
            "seed": 42,
            "highvram": True,
            "save_every_n_epochs": 1,
            "full_bf16": True,
            "save_precision": "bf16",
            "sdpa": True,
            "model_prediction_type": "raw",
            "discrete_flow_shift": 3.1582,
            "disable_mmap_load_safetensors": True,
        }
        for k, v in fixed_params.items():
            if k not in combined_config:
                combined_config[k] = v
        
        return (combined_config,)


class Flux16GBParams:
    """FLUX 16GB显存优化参数节点 - 基于train_flux_16gb.sh脚本配置"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        """定义输入类型"""
        return {
            "required": {
                "pretrained_model_name_or_path": (folder_paths.get_filename_list("checkpoints"), {
                    "default": "", 
                    "multiline": False, 
                    "tooltip": "预训练模型路径 - 选择要训练的基础模型。具体影响：SD1.5(通用模型，兼容性好，训练稳定) → SDXL(高质量模型，细节丰富，需要更多显存) → SDXL Turbo(快速模型，训练快但效果可能不如SDXL)。选择正确的模型对训练效果至关重要。"
                }),
                # 基础训练参数
                "train_data_dir": ("STRING", {
                    "default": "", 
                    "multiline": False, 
                    "tooltip": "训练数据目录路径 - 包含训练图像的文件夹路径。具体要求：文件夹内应包含子文件夹，每个子文件夹包含图像文件和对应的文本描述文件(.txt或.caption)。例如：/path/to/train_data/character1/image1.jpg + image1.txt。注意：train_data_dir应该是包含子文件夹的父目录，而不是直接包含图像的目录。"
                }),
                "output_dir": ("STRING", {
                    "default": "./output", 
                    "multiline": False, 
                    "tooltip": "模型输出目录 - 训练完成的LoRA模型保存路径。具体影响：相对路径(相对于ComfyUI根目录) → 绝对路径(完整路径，推荐) → 默认路径(自动生成)。建议使用绝对路径，避免路径问题。"
                }),
                "output_name": ("STRING", {
                    "default": "trained_flux_lora", 
                    "multiline": False, 
                    "tooltip": "输出模型名称 - 训练完成的LoRA模型文件名(不含扩展名)。具体影响：简单名称(如'my_lora') → 描述性名称(如'character_style_v1') → 版本化名称(如'portrait_style_v2.1')。建议使用描述性名称，便于管理多个模型。"
                }),
                # 模型类型配置
                "config_file": ("STRING", {
                    "default": "", 
                    "multiline": False, 
                    "tooltip": "TOML配置文件路径 - 可选的配置文件路径，用于加载预定义的训练配置。具体用途：空值(使用节点参数) → TOML文件路径(加载配置文件，适合批量训练或复杂配置)。配置文件可以包含完整的训练参数设置，便于重复使用和版本管理。"
                }),
                # 基础训练参数 - 与脚本完全一致
                "network_dim": ("INT", {
                    "default": 32,  # 更新为更好的默认值
                    "min": 1, 
                    "max": 64,
                    "tooltip": "LoRA网络维度，推荐值32(优秀效果，适合16GB显存)"
                }),
                "network_alpha": ("FLOAT", {
                    "default": 16.0,  # 更新为更好的默认值
                    "min": 0.1, 
                    "max": 128.0,  # 修复：从10.0改为128.0，与其他配置保持一致
                    "tooltip": "LoRA网络alpha值，推荐为network_dim的1/2"
                }),
                "learning_rate": ("FLOAT", {
                    "default": 1e-4,  # 脚本中的默认值
                    "min": 1e-6, 
                    "max": 1e-2,
                    "tooltip": "学习率，脚本默认值为1e-4"
                }),
                "max_train_steps": ("INT", {
                    "default": 1200,  # 更新为更好的默认值
                    "min": 1, 
                    "max": 10000,
                    "tooltip": "最大训练步数，推荐值1200(深入学习，效果优秀)"
                }),
                # 新增：训练轮数控制（优先级高于步数）
                "max_train_epochs": ("INT", {
                    "default": 5,  # 文档中的默认值
                    "min": 1, 
                    "max": 100,
                    "tooltip": "最大训练轮数，优先级高于max_train_steps，推荐值5轮"
                }),
                # 新增：分离学习率（重要参数）
                "unet_lr": ("FLOAT", {
                    "default": 8e-4,  # 文档中的默认值
                    "min": 1e-6, 
                    "max": 1e-2,
                    "tooltip": "UNet学习率，比基础学习率高，提升图像生成质量"
                }),
                "text_encoder_lr": ("FLOAT", {
                    "default": 1e-5,  # 文档中的默认值
                    "min": 1e-6, 
                    "max": 1e-3,
                    "tooltip": "文本编码器学习率，比基础学习率低，保护文本理解能力"
                }),
                 # 保存与中断
                "save_every_n_steps": ("INT", {
                    "default": 500,
                    "min": 100,
                    "max": 5000,
                    "tooltip": "每多少步保存一次模型"
                }),
                "train_batch_size": ("INT", {
                    "default": 1,  # 脚本中的默认值
                    "min": 1, 
                    "max": 4,
                    "tooltip": "训练批次大小，脚本默认值为1"
                }),
                
                # 分辨率参数 - 与脚本完全一致
                "resolution": ("INT", {
                    "default": 512,  # 脚本中的默认值
                    "min": 256,
                    "max": 1024,
                    "tooltip": "训练图像分辨率，脚本默认值为512"
                }),
                
                # 必需的模型路径参数
                "clip_l": (["None"] + folder_paths.get_filename_list("clip"), {
                    "default": "None",
                    "tooltip": "CLIP-L模型文件，选择None时使用基础模型内置CLIP-L"
                }),
                "t5xxl": (["None"] + folder_paths.get_filename_list("text_encoders"), {
                    "default": "None",
                    "tooltip": "T5XXL模型文件，选择None时使用基础模型内置T5XXL"
                }),
                "ae": (["None"] + folder_paths.get_filename_list("vae"), {
                    "default": "None",
                    "tooltip": "AutoEncoder模型文件，选择None时使用基础模型内置AE"
                }),
                
                # FLUX特定参数 - 与脚本完全一致
                "guidance_scale": ("FLOAT", {
                    "default": 1.0,  # 脚本中的默认值
                    "min": 0.1, 
                    "max": 10.0,
                    "tooltip": "引导尺度，脚本默认值为1.0"
                }),
                "timestep_sampling": (["shift", "uniform"], {
                    "default": "shift",  # 脚本中的默认值
                    "tooltip": "时间步采样方式，脚本默认值为shift"
                }),
                "model_prediction_type": (["raw", "v_prediction", "epsilon"], {
                    "default": "raw",  # 脚本中的默认值
                    "tooltip": "模型预测类型，脚本默认值为raw"
                }),
                "discrete_flow_shift": ("FLOAT", {
                    "default": 3.1582,  # 脚本中的默认值
                    "min": 0.1, 
                    "max": 10.0,
                    "tooltip": "离散流偏移，脚本默认值为3.1582"
                }),
                
                # 训练控制参数 - 与脚本完全一致
                "network_train_unet_only": ("BOOLEAN", {
                    "default": True,  # 脚本中的默认值
                    "tooltip": "仅训练UNet，脚本默认值为True"
                }),
                
                # 内存优化参数 - 与脚本完全一致
                "gradient_accumulation_steps": ("INT", {
                    "default": 1,  # 脚本中未指定，使用默认值
                    "min": 1,
                    "max": 16,
                    "tooltip": "梯度累积步数"
                }),
                "max_data_loader_n_workers": ("INT", {
                    "default": 8,  # 从4增加到8，进一步提高数据加载速度
                    "min": 0,
                    "max": 8,  # 增加最大值限制
                    "tooltip": "数据加载器工作进程数，推荐值8(高速数据加载)"
                }),
                
                # 混合精度设置 - 与脚本完全一致
                "mixed_precision": (["no", "fp16", "bf16"], {
                    "default": "bf16",  # 脚本中的默认值
                    "tooltip": "混合精度训练，脚本默认值为bf16"
                }),
                
                # 优化器设置 - 与脚本完全一致
                "optimizer_type": (["AdamW", "AdamW8bit", "Lion", "SGDNesterov", "DAdaptation", "Prodigy", "AdaFactor"], {
                    "default": "Lion",  # 从AdamW8bit改为Lion，最快的优化器
                    "tooltip": "优化器类型，推荐Lion(速度最快)"
                }),
                
                # 缓存设置 - 与脚本完全一致
                "cache_latents": ("BOOLEAN", {
                    "default": True,  # 脚本中未指定，使用默认值
                    "tooltip": "缓存潜在表示"
                }),
                "cache_latents_to_disk": ("BOOLEAN", {
                    "default": False,  # 脚本中未指定，使用默认值
                    "tooltip": "将潜在表示缓存到磁盘"
                }),
                "cache_text_encoder_outputs": ("BOOLEAN", {
                    "default": True,  # 脚本中的默认值
                    "tooltip": "缓存文本编码器输出，脚本默认值为True"
                }),
                "cache_text_encoder_outputs_to_disk": ("BOOLEAN", {
                    "default": True,  # 脚本中的默认值
                    "tooltip": "将文本编码器输出缓存到磁盘，脚本默认值为True"
                }),
                
                # 高级内存优化 - 与脚本完全一致
                "blocks_to_swap": ("INT", {
                    "default": 8,  # 从12减少到8，提高训练速度
                    "min": 0,
                    "max": 20,
                    "tooltip": "要交换的块数量，推荐值8(平衡速度和显存)"
                }),
                "cpu_offload_checkpointing": ("BOOLEAN", {
                    "default": False,  # 脚本中未指定，使用默认值
                    "tooltip": "启用CPU卸载检查点"
                }),
                "fp8_base": ("BOOLEAN", {
                    "default": True,  # 脚本中的默认值
                    "tooltip": "使用FP8基础模型，脚本默认值为True"
                }),
                "disable_mmap_load_safetensors": ("BOOLEAN", {
                    "default": True,  # 脚本中的默认值
                    "tooltip": "禁用内存映射加载，脚本默认值为True"
                }),
                
                # 新增：重要训练效果参数
                "keep_tokens": ("INT", {
                    "default": 1,  # 文档中的默认值
                    "min": 0, 
                    "max": 10,
                    "tooltip": "保留标记数，保护重要概念token不被dropout影响"
                }),
                "lr_scheduler_num_cycles": ("INT", {
                    "default": 2,  # 文档中的默认值
                    "min": 1, 
                    "max": 10,
                    "tooltip": "学习率调度器周期数，2个周期可能跳出局部最优"
                }),
                "lr_decay_steps": ("FLOAT", {
                    "default": 0.5,  # 文档中的默认值
                    "min": 0.1, 
                    "max": 1.0,
                    "tooltip": "学习率衰减步数比例，0.5表示在50%步数时开始衰减"
                }),
                "lr_scheduler_min_lr_ratio": ("FLOAT", {
                    "default": 0.1,  # 文档中的默认值
                    "min": 0.01, 
                    "max": 1.0,
                    "tooltip": "最小学习率比例，避免学习率过低导致训练停滞"
                }),
                "caption_extension": ("STRING", {
                    "default": ".txt",  # 文档中的默认值
                    "tooltip": "标注文件扩展名，确保正确读取标注文件"
                }),
                "vae_batch_size": ("INT", {
                    "default": 4,  # 文档中的默认值
                    "min": 1, 
                    "max": 8,
                    "tooltip": "VAE批处理大小，提高VAE处理效率"
                }),
                "apply_t5_attn_mask": ("BOOLEAN", {
                    "default": True,  # 文档中的默认值
                    "tooltip": "应用T5注意力掩码，提升T5模型文本理解能力"
                }),
                "sigmoid_scale": ("FLOAT", {
                    "default": 1.0,  # 文档中的默认值
                    "min": 0.1, 
                    "max": 10.0,
                    "tooltip": "Sigmoid缩放，控制激活函数强度"
                }),
                "bucket_no_upscale": ("BOOLEAN", {
                    "default": True,  # 文档中的默认值
                    "tooltip": "分桶不上采样，保持图像原始比例避免失真"
                }),
                
                # 额外参数 - 与脚本完全一致
                "lr_scheduler": (["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup", "cosine_with_min_lr"], {
                    "default": "cosine_with_min_lr",  # 更新为文档中的默认值
                    "tooltip": "学习率调度器，文档默认值为cosine_with_min_lr"
                }),
                "max_grad_norm": ("FLOAT", {
                    "default": 0.0,  # 脚本中的默认值
                    "min": 0.0,
                    "max": 10.0,
                    "tooltip": "最大梯度范数，脚本默认值为0.0"
                }),
                "enable_bucket": ("BOOLEAN", {
                    "default": True,  # 脚本中的默认值
                    "tooltip": "启用bucket优化，脚本默认值为True"
                }),
                "min_bucket_reso": ("INT", {
                    "default": 256,  # 脚本中的默认值
                    "min": 128,
                    "max": 512,
                    "tooltip": "最小bucket分辨率，脚本默认值为256"
                }),
                "max_bucket_reso": ("INT", {
                    "default": 512,  # 脚本中的默认值
                    "min": 256,
                    "max": 1024,
                    "tooltip": "最大bucket分辨率，脚本默认值为512"
                }),
                "bucket_reso_steps": ("INT", {
                    "default": 64,  # 脚本中的默认值
                    "min": 32,
                    "max": 128,
                    "tooltip": "bucket分辨率步长，脚本默认值为64"
                }),
            }
        }
    
    RETURN_TYPES = ("FLUX_TRAINING_PARAMS",)
    FUNCTION = "create_16gb_params"
    CATEGORY = FLUX_NEW_PARAMS_CATEGORY
    

    def validate_inputs(kwargs):
        """验证输入参数 - 返回True或错误字符串"""
 

        # 验证训练数据目录 - 只有当用户输入了值时才验证
        train_data_dir = kwargs.get("train_data_dir", "")
        if train_data_dir:  # 只有当用户输入了值时才验证
            if not os.path.exists(train_data_dir):
                return f"错误：训练数据目录不存在: {train_data_dir}"
            
            if not os.path.isdir(train_data_dir):
                return f"错误：训练数据路径不是目录: {train_data_dir}"
            
            # 检查是否有图片文件
            image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.webp']
            has_images = False
            try:
                for root, dirs, files in os.walk(train_data_dir):
                    for file in files:
                        if any(file.lower().endswith(ext) for ext in image_extensions):
                            has_images = True
                            break
                    if has_images:
                        break
                
                if not has_images:
                    return f"错误：训练数据目录中未找到图片文件: {train_data_dir}"
            except Exception as e:
                return f"错误：无法访问训练数据目录: {e}"
        
        # 验证输出目录 - 只有当用户输入了值时才验证
        output_dir = kwargs.get("output_dir", "")
        if output_dir:  # 只有当用户输入了值时才验证
            # 尝试创建输出目录
            try:
                os.makedirs(output_dir, exist_ok=True)
            except Exception as e:
                return f"错误：无法创建输出目录 {output_dir}: {e}"
        
        # 验证输出名称 - 只有当用户输入了值时才验证
        output_name = kwargs.get("output_name", "")
        if output_name:  # 只有当用户输入了值时才验证
            # 检查输出名称是否包含非法字符
            import re
            if re.search(r'[<>:"/\\|?*]', output_name):
                return f"错误：输出名称包含非法字符: {output_name}"
        
        # 验证配置文件路径（如果提供）
        config_file = kwargs.get("config_file", "")
        if config_file:
            if not os.path.exists(config_file):
                return f"错误：配置文件不存在: {config_file}"
            
            if not config_file.lower().endswith('.toml'):
                return f"错误：配置文件必须是TOML格式: {config_file}"
        
        # 验证必需模型文件 - 只有当用户选择了值时才验证
        clip_l = kwargs.get("clip_l", "")
        t5xxl = kwargs.get("t5xxl", "")
        ae = kwargs.get("ae", "")
        
        if clip_l and not os.path.exists(folder_paths.get_full_path("clip", clip_l)):
            return f"错误：CLIP-L模型文件不存在: {clip_l}"
        
        if t5xxl and not os.path.exists(folder_paths.get_full_path("text_encoders", t5xxl)):
            return f"错误：T5XXL模型文件不存在: {t5xxl}"
        
        if ae and not os.path.exists(folder_paths.get_full_path("vae", ae)):
            return f"错误：AutoEncoder模型文件不存在: {ae}"
        
        return True  # 验证通过
    def create_16gb_params(self,**kwargs):
        """创建FLUX 16GB显存优化参数 - 与train_flux_16gb.sh脚本完全一致"""
        
        # 处理"None"值，删除对应参数
        none_params = ["clip_l", "t5xxl", "ae"]
        for param in none_params:
            if kwargs.get(param) == "None":
                kwargs.pop(param, None)
                print(f"[INFO] 参数 {param} 设置为None，已从训练参数中移除")
        
        # 处理分辨率格式
        kwargs['resolution'] = f"{kwargs['resolution']},{kwargs['resolution']}"
        
        # 添加一些与训练效果相关的硬编码参数
        kwargs.update({
            "model_prediction_type": "raw",
            "discrete_flow_shift": 3.1582,
        })

         # 将相对路径转换为绝对路径
        for key in ["clip_l", "t5xxl", "ae","pretrained_model_name_or_path"]:
            if kwargs[key]:
                try:
                    if key == "clip_l":
                        kwargs[key] = folder_paths.get_full_path("clip", kwargs[key])
                    elif key == "t5xxl":
                        kwargs[key] = folder_paths.get_full_path("text_encoders", kwargs[key])
                    elif key == "ae":
                        kwargs[key] = folder_paths.get_full_path("vae", kwargs[key])
                    elif key == "pretrained_model_name_or_path":
                        kwargs[key] = folder_paths.get_full_path("checkpoints", kwargs[key])
                except Exception as e:
                    logger.warning(f"无法获取 {key} 的完整路径: {e}")

        # 处理训练轮数控制（优先级高于步数）
        if kwargs.get("max_train_epochs", 0) > 0:
            # 如果设置了训练轮数，则清除训练步数设置
            kwargs.pop("max_train_steps", None)
            logger.info(f"使用训练轮数控制: {kwargs['max_train_epochs']} 轮")
        else:
            # 如果没有设置训练轮数，则使用训练步数
            kwargs.pop("max_train_epochs", None)
            logger.info(f"使用训练步数控制: {kwargs.get('max_train_steps', 1200)} 步")

        # 处理分离学习率
        if kwargs.get("unet_lr") and kwargs.get("text_encoder_lr"):
            # 如果设置了分离学习率，则使用分离学习率
            logger.info(f"使用分离学习率: UNet={kwargs['unet_lr']}, TextEncoder={kwargs['text_encoder_lr']}")
        else:
            # 如果没有设置分离学习率，则使用统一学习率
            logger.info(f"使用统一学习率: {kwargs.get('learning_rate', 1e-4)}")

        # 根据优化器类型设置optimizer_args
        optimizer_type = kwargs.get("optimizer_type", "AdamW")
        if optimizer_type == "AdaFactor":
            kwargs["optimizer_args"] = ["relative_step=False", "scale_parameter=False", "warmup_init=False"]
        elif optimizer_type == "AdamW":
            kwargs["optimizer_args"] = ["weight_decay=0.01", "betas=0.9,0.999"]
        elif optimizer_type == "AdamW8bit":
            kwargs["optimizer_args"] = ["weight_decay=0.01", "betas=0.9,0.999"]
        elif optimizer_type == "Lion":
            kwargs["optimizer_args"] = ["weight_decay=0.01"]
        elif optimizer_type == "Lion8bit":
            kwargs["optimizer_args"] = ["weight_decay=0.01"]
        elif optimizer_type == "SGDNesterov":
            kwargs["optimizer_args"] = ["momentum=0.9", "weight_decay=0.01"]
        elif optimizer_type == "SGDNesterov8bit":
            kwargs["optimizer_args"] = ["momentum=0.9", "weight_decay=0.01"]
        else:
            kwargs["optimizer_args"] = ["weight_decay=0.01", "betas=0.9,0.999"]
        
        # 脚本中的固定参数
        kwargs.update({
            "save_model_as": "safetensors",
            "sdpa": True,
            "persistent_data_loader_workers": True,
            "seed": 42,
            "gradient_checkpointing": True,
            "save_precision": "bf16",
            "network_module": "networks.lora_flux",
            "highvram": True,
            "save_every_n_epochs": 1,
            "full_bf16": True,
        })
        
        # 记录重要参数信息
        logger.info(f"FLUX 16GB参数已创建（包含训练效果优化参数）:")
        logger.info(f"  网络维度: {kwargs['network_dim']}")
        logger.info(f"  网络Alpha: {kwargs['network_alpha']}")
        logger.info(f"  分辨率: {kwargs['resolution']}")
        logger.info(f"  批次大小: {kwargs['train_batch_size']}")
        logger.info(f"  基础学习率: {kwargs.get('learning_rate', 'N/A')}")
        if kwargs.get("unet_lr"):
            logger.info(f"  UNet学习率: {kwargs['unet_lr']}")
        if kwargs.get("text_encoder_lr"):
            logger.info(f"  文本编码器学习率: {kwargs['text_encoder_lr']}")
        if kwargs.get("max_train_epochs"):
            logger.info(f"  训练轮数: {kwargs['max_train_epochs']}")
        if kwargs.get("max_train_steps"):
            logger.info(f"  训练步数: {kwargs['max_train_steps']}")
        logger.info(f"  混合精度: {kwargs['mixed_precision']}")
        logger.info(f"  优化器: {kwargs['optimizer_type']}")
        logger.info(f"  学习率调度器: {kwargs.get('lr_scheduler', 'N/A')}")
        logger.info(f"  学习率周期数: {kwargs.get('lr_scheduler_num_cycles', 'N/A')}")
        logger.info(f"  T5注意力掩码: {kwargs.get('apply_t5_attn_mask', 'N/A')}")
        logger.info(f"  保留标记数: {kwargs.get('keep_tokens', 'N/A')}")
        logger.info(f"  块交换: {kwargs['blocks_to_swap']}")
        logger.info(f"  FP8基础: {kwargs['fp8_base']}")
        
        return (kwargs,)

# 节点注册
NODE_CLASS_MAPPINGS = {
    "FluxDependencyConfig": FluxDependencyConfig,
    "FluxTrainingConfig": FluxTrainingConfig,
    "FluxPerformanceConfig": FluxPerformanceConfig,
    "FluxConfigCombiner": FluxConfigCombiner,
    "Flux16GBParams": Flux16GBParams,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "FluxDependencyConfig": "FLUX依赖配置",
    "FluxTrainingConfig": "FLUX训练配置", 
    "FluxPerformanceConfig": "FLUX性能配置",
    "FluxConfigCombiner": "FLUX配置合并",
    "Flux16GBParams": "FLUX 16GB显存优化参数",
}